import sys
from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QLabel, 
    QLineEdit, QPushButton, QFileDialog, QComboBox, QSpinBox, QCheckBox,
    QMessageBox, QGroupBox, QColorDialog
)
from PyQt6.QtGui import QColor, QPixmap, QIcon
from PyQt6.QtCore import Qt
from PIL import Image, ImageDraw
import os

class PhotoGeneratorApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("证件照生成工具")
        self.setWindowIcon(QIcon("icon.png"))  # 替换为您的图标路径
        self.setGeometry(100, 100, 800, 600)
        
        self.initUI()
        self.setupConnections()
        
        # 默认值
        self.photo_sizes = {
            "1寸": (295, 413),
            "2寸": (413, 579),
            "小1寸": (260, 378),
            "小2寸": (413, 531),
            "5寸": (1200, 840),
            "6寸": (1200, 1800),
            "自定义": None
        }
        
        self.paper_sizes = {
            "5寸": (1200, 840),
            "6寸": (1200, 1800),
            "A4": (2480, 3508),
            "自定义": None
        }
        
        self.bg_colors = {
            "蓝色": (0, 191, 255),
            "红色": (255, 0, 0),
            "白色": (255, 255, 255),
            "自定义": None
        }
        
        # 初始化UI值
        self.photoSizeCombo.addItems(self.photo_sizes.keys())
        self.paperSizeCombo.addItems(self.paper_sizes.keys())
        self.bgColorCombo.addItems(self.bg_colors.keys())
        
        self.widthSpin.setValue(295)
        self.heightSpin.setValue(413)
        self.paperWidthSpin.setValue(1200)
        self.paperHeightSpin.setValue(1800)
        self.marginSpin.setValue(50)
        
        self.bgColorPreview.setStyleSheet("background-color: rgb(0, 191, 255)")
        
    def initUI(self):
        main_widget = QWidget()
        main_layout = QVBoxLayout()
        
        # 输入设置区域
        input_group = QGroupBox("输入设置")
        input_layout = QVBoxLayout()
        
        # 选择输入文件或文件夹
        file_layout = QHBoxLayout()
        self.inputFileRadio = QCheckBox("单张照片")
        self.inputFileRadio.setChecked(True)
        self.inputFolderRadio = QCheckBox("批量处理")
        file_layout.addWidget(QLabel("输入方式:"))
        file_layout.addWidget(self.inputFileRadio)
        file_layout.addWidget(self.inputFolderRadio)
        file_layout.addStretch()
        
        self.inputPathEdit = QLineEdit()
        self.inputPathEdit.setPlaceholderText("请选择图片文件或文件夹")
        self.browseButton = QPushButton("浏览...")
        
        input_layout.addLayout(file_layout)
        input_layout.addWidget(QLabel("输入路径:"))
        input_layout.addWidget(self.inputPathEdit)
        input_layout.addWidget(self.browseButton)
        input_group.setLayout(input_layout)
        
        # 输出设置区域
        output_group = QGroupBox("输出设置")
        output_layout = QVBoxLayout()
        
        self.outputPathEdit = QLineEdit()
        self.outputPathEdit.setPlaceholderText("请选择输出文件夹")
        self.outputBrowseButton = QPushButton("浏览...")
        
        output_layout.addWidget(QLabel("输出文件夹:"))
        output_layout.addWidget(self.outputPathEdit)
        output_layout.addWidget(self.outputBrowseButton)
        output_group.setLayout(output_layout)
        
        # 证件照设置区域
        photo_group = QGroupBox("证件照设置")
        photo_layout = QVBoxLayout()
        
        # 证件照尺寸
        size_layout = QHBoxLayout()
        self.photoSizeCombo = QComboBox()
        self.widthSpin = QSpinBox()
        self.widthSpin.setRange(10, 9999)
        self.heightSpin = QSpinBox()
        self.heightSpin.setRange(10, 9999)
        size_layout.addWidget(QLabel("标准尺寸:"))
        size_layout.addWidget(self.photoSizeCombo)
        size_layout.addWidget(QLabel("宽度:"))
        size_layout.addWidget(self.widthSpin)
        size_layout.addWidget(QLabel("高度:"))
        size_layout.addWidget(self.heightSpin)
        
        # 背景颜色
        color_layout = QHBoxLayout()
        self.bgColorCombo = QComboBox()
        self.bgColorPreview = QLabel()
        self.bgColorPreview.setFixedSize(50, 25)
        self.customColorButton = QPushButton("自定义颜色")
        color_layout.addWidget(QLabel("背景颜色:"))
        color_layout.addWidget(self.bgColorCombo)
        color_layout.addWidget(self.bgColorPreview)
        color_layout.addWidget(self.customColorButton)
        color_layout.addStretch()
        
        photo_layout.addLayout(size_layout)
        photo_layout.addLayout(color_layout)
        photo_group.setLayout(photo_layout)
        
        # 排版设置区域
        arrange_group = QGroupBox("排版设置")
        arrange_layout = QVBoxLayout()
        
        # 相纸尺寸
        paper_size_layout = QHBoxLayout()
        self.paperSizeCombo = QComboBox()
        self.paperWidthSpin = QSpinBox()
        self.paperWidthSpin.setRange(10, 9999)
        self.paperHeightSpin = QSpinBox()
        self.paperHeightSpin.setRange(10, 9999)
        paper_size_layout.addWidget(QLabel("相纸尺寸:"))
        paper_size_layout.addWidget(self.paperSizeCombo)
        paper_size_layout.addWidget(QLabel("宽度:"))
        paper_size_layout.addWidget(self.paperWidthSpin)
        paper_size_layout.addWidget(QLabel("高度:"))
        paper_size_layout.addWidget(self.paperHeightSpin)
        
        # 边距设置
        margin_layout = QHBoxLayout()
        self.marginSpin = QSpinBox()
        self.marginSpin.setRange(0, 500)
        margin_layout.addWidget(QLabel("边距(像素):"))
        margin_layout.addWidget(self.marginSpin)
        margin_layout.addStretch()
        
        # 输出选项
        options_layout = QHBoxLayout()
        self.saveSingleCheck = QCheckBox("保存单张证件照")
        self.saveSingleCheck.setChecked(True)
        self.saveArrangedCheck = QCheckBox("保存排版照片")
        self.saveArrangedCheck.setChecked(True)
        options_layout.addWidget(self.saveSingleCheck)
        options_layout.addWidget(self.saveArrangedCheck)
        options_layout.addStretch()
        
        arrange_layout.addLayout(paper_size_layout)
        arrange_layout.addLayout(margin_layout)
        arrange_layout.addLayout(options_layout)
        arrange_group.setLayout(arrange_layout)
        
        # 操作按钮
        button_layout = QHBoxLayout()
        self.generateButton = QPushButton("生成证件照")
        self.previewButton = QPushButton("预览")
        button_layout.addWidget(self.generateButton)
        button_layout.addWidget(self.previewButton)
        
        # 预览区域
        self.previewLabel = QLabel()
        self.previewLabel.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.previewLabel.setFixedHeight(200)
        self.previewLabel.setStyleSheet("border: 1px solid gray;")
        
        # 添加到主布局
        main_layout.addWidget(input_group)
        main_layout.addWidget(output_group)
        main_layout.addWidget(photo_group)
        main_layout.addWidget(arrange_group)
        main_layout.addLayout(button_layout)
        main_layout.addWidget(self.previewLabel)
        main_layout.addStretch()
        
        main_widget.setLayout(main_layout)
        self.setCentralWidget(main_widget)
    
    def setupConnections(self):
        self.browseButton.clicked.connect(self.browseInput)
        self.outputBrowseButton.clicked.connect(self.browseOutput)
        self.generateButton.clicked.connect(self.generatePhotos)
        self.previewButton.clicked.connect(self.previewPhoto)
        self.customColorButton.clicked.connect(self.chooseCustomColor)
        
        self.photoSizeCombo.currentTextChanged.connect(self.updatePhotoSize)
        self.paperSizeCombo.currentTextChanged.connect(self.updatePaperSize)
        self.bgColorCombo.currentTextChanged.connect(self.updateBgColor)
        
        self.widthSpin.valueChanged.connect(self.updateCustomPhotoSize)
        self.heightSpin.valueChanged.connect(self.updateCustomPhotoSize)
        self.paperWidthSpin.valueChanged.connect(self.updateCustomPaperSize)
        self.paperHeightSpin.valueChanged.connect(self.updateCustomPaperSize)
    
    def browseInput(self):
        if self.inputFileRadio.isChecked():
            file_path, _ = QFileDialog.getOpenFileName(
                self, "选择照片文件", "", "图片文件 (*.jpg *.jpeg *.png)"
            )
            if file_path:
                self.inputPathEdit.setText(file_path)
                # 自动预览
                self.previewPhoto()
        else:
            folder_path = QFileDialog.getExistingDirectory(
                self, "选择照片文件夹"
            )
            if folder_path:
                self.inputPathEdit.setText(folder_path)
    
    def browseOutput(self):
        folder_path = QFileDialog.getExistingDirectory(
            self, "选择输出文件夹"
        )
        if folder_path:
            self.outputPathEdit.setText(folder_path)
    
    def chooseCustomColor(self):
        color = QColorDialog.getColor()
        if color.isValid():
            self.bg_colors["自定义"] = (color.red(), color.green(), color.blue())
            self.bgColorCombo.setCurrentText("自定义")
            self.updateBgColorPreview(color.red(), color.green(), color.blue())
    
    def updatePhotoSize(self, text):
        if text != "自定义" and self.photo_sizes[text]:
            width, height = self.photo_sizes[text]
            self.widthSpin.setValue(width)
            self.heightSpin.setValue(height)
    
    def updatePaperSize(self, text):
        if text != "自定义" and self.paper_sizes[text]:
            width, height = self.paper_sizes[text]
            self.paperWidthSpin.setValue(width)
            self.paperHeightSpin.setValue(height)
    
    def updateBgColor(self, text):
        if text != "自定义" and self.bg_colors[text]:
            r, g, b = self.bg_colors[text]
            self.updateBgColorPreview(r, g, b)
    
    def updateBgColorPreview(self, r, g, b):
        self.bgColorPreview.setStyleSheet(f"background-color: rgb({r}, {g}, {b})")
    
    def updateCustomPhotoSize(self):
        self.photoSizeCombo.setCurrentText("自定义")
        self.photo_sizes["自定义"] = (self.widthSpin.value(), self.heightSpin.value())
    
    def updateCustomPaperSize(self):
        self.paperSizeCombo.setCurrentText("自定义")
        self.paper_sizes["自定义"] = (self.paperWidthSpin.value(), self.paperHeightSpin.value())
    
    def getCurrentSettings(self):
        settings = {
            "input_path": self.inputPathEdit.text(),
            "output_folder": self.outputPathEdit.text(),
            "photo_size": (self.widthSpin.value(), self.heightSpin.value()),
            "paper_size": (self.paperWidthSpin.value(), self.paperHeightSpin.value()),
            "margin": self.marginSpin.value(),
            "save_single": self.saveSingleCheck.isChecked(),
            "save_arranged": self.saveArrangedCheck.isChecked(),
            "is_batch": self.inputFolderRadio.isChecked()
        }
        
        # 获取背景颜色
        bg_color_text = self.bgColorCombo.currentText()
        if bg_color_text == "自定义" and self.bg_colors["自定义"]:
            settings["bg_color"] = self.bg_colors["自定义"]
        else:
            settings["bg_color"] = self.bg_colors.get(bg_color_text, (0, 191, 255))
        
        return settings
    
    def validateSettings(self, settings):
        if not settings["input_path"]:
            QMessageBox.warning(self, "警告", "请选择输入文件或文件夹!")
            return False
        
        if not settings["output_folder"]:
            QMessageBox.warning(self, "警告", "请选择输出文件夹!")
            return False
        
        if not os.path.exists(settings["input_path"]):
            QMessageBox.warning(self, "警告", "输入路径不存在!")
            return False
        
        if settings["is_batch"] and not os.path.isdir(settings["input_path"]):
            QMessageBox.warning(self, "警告", "批量处理需要选择文件夹!")
            return False
        
        if not settings["is_batch"] and not os.path.isfile(settings["input_path"]):
            QMessageBox.warning(self, "警告", "单张处理需要选择文件!")
            return False
        
        return True
    
    def previewPhoto(self):
        settings = self.getCurrentSettings()
        
        if not settings["input_path"] or not os.path.isfile(settings["input_path"]):
            return
        
        try:
            # 打开图像
            image = Image.open(settings["input_path"])
            
            # 调整图像尺寸
            resized_image = self.resize_image(image, settings["photo_size"])
            
            # 替换背景颜色
            final_image = self.change_background_color(resized_image, settings["bg_color"])
            
            # 转换为QPixmap显示
            final_image.save("preview_temp.jpg")  # 临时保存
            pixmap = QPixmap("preview_temp.jpg")
            os.remove("preview_temp.jpg")  # 删除临时文件
            
            # 缩放以适应预览区域
            scaled_pixmap = pixmap.scaled(
                self.previewLabel.width(), self.previewLabel.height(),
                Qt.AspectRatioMode.KeepAspectRatio,
                Qt.TransformationMode.SmoothTransformation
            )
            self.previewLabel.setPixmap(scaled_pixmap)
            
        except Exception as e:
            QMessageBox.warning(self, "预览错误", f"预览时出错: {str(e)}")
    
    def generatePhotos(self):
        settings = self.getCurrentSettings()
        
        if not self.validateSettings(settings):
            return
        
        try:
            if settings["is_batch"]:
                self.batch_process_images(
                    settings["input_path"], settings["output_folder"],
                    settings["photo_size"], settings["bg_color"],
                    settings["paper_size"], settings["margin"],
                    settings["save_single"], settings["save_arranged"]
                )
            else:
                self.make_id_photo(
                    settings["input_path"], settings["output_folder"],
                    settings["photo_size"], settings["bg_color"],
                    settings["paper_size"], settings["margin"],
                    settings["save_single"], settings["save_arranged"]
                )
            
            QMessageBox.information(self, "完成", "证件照生成完成!")
        
        except Exception as e:
            QMessageBox.critical(self, "错误", f"生成证件照时出错: {str(e)}")
    
    # 以下是原功能函数的移植，稍作修改以适应PyQt界面
    
    def resize_image(self, image, size):
        """调整图像尺寸"""
        return image.resize(size, Image.Resampling.LANCZOS)
    
    def change_background_color(self, image, bg_color):
        """替换图像背景颜色"""
        # 创建一个新图像，填充背景颜色
        background = Image.new("RGB", image.size, bg_color)

        # 如果图像有透明通道（如 PNG），将透明部分替换为背景颜色
        if image.mode in ("RGBA", "LA"):
            background.paste(image, mask=image.split()[-1])  # 使用透明通道作为掩码
        else:
            background.paste(image)

        return background
    
    def add_crop_lines(self, draw, x, y, photo_size, paper_size, margin=50, line_width=1):
        """在证件照边缘添加裁剪线，并延伸到相纸边缘"""
        # 裁剪线颜色（黑色）
        crop_color = (0, 0, 0)

        # 上边缘裁剪线（从左到右）
        draw.line([(x, y), (x + photo_size[0], y)], fill=crop_color, width=line_width)

        # 下边缘裁剪线（从左到右）
        draw.line([(x, y + photo_size[1]), (x + photo_size[0], y + photo_size[1])], fill=crop_color, width=line_width)

        # 左边缘裁剪线（从上到下）
        draw.line([(x, y), (x, y + photo_size[1])], fill=crop_color, width=line_width)

        # 右边缘裁剪线（从上到下）
        draw.line([(x + photo_size[0], y), (x + photo_size[0], y + photo_size[1])], fill=crop_color, width=line_width)

        # 延伸到相纸边缘的裁剪线
        # 上边缘延伸到相纸顶部
        draw.line([(x, y), (x, 0)], fill=crop_color, width=line_width)
        draw.line([(x + photo_size[0], y), (x + photo_size[0], 0)], fill=crop_color, width=line_width)

        # 下边缘延伸到相纸底部
        draw.line([(x, y + photo_size[1]), (x, paper_size[1])], fill=crop_color, width=line_width)
        draw.line([(x + photo_size[0], y + photo_size[1]), (x + photo_size[0], paper_size[1])], fill=crop_color, width=line_width)

        # 左边缘延伸到相纸左侧
        draw.line([(x, y), (0, y)], fill=crop_color, width=line_width)
        draw.line([(x, y + photo_size[1]), (0, y + photo_size[1])], fill=crop_color, width=line_width)

        # 右边缘延伸到相纸右侧
        draw.line([(x + photo_size[0], y), (paper_size[0], y)], fill=crop_color, width=line_width)
        draw.line([(x + photo_size[0], y + photo_size[1]), (paper_size[0], y + photo_size[1])], fill=crop_color, width=line_width)
    
    def arrange_photos_on_paper(self, photo, paper_size, photo_size, margin=50):
        """将多张证件照均匀排列在相纸上，并自动调整方向以充分利用相纸"""
        paper = Image.new("RGB", paper_size, (255, 255, 255))  # 创建白色背景的相纸
        draw = ImageDraw.Draw(paper)

        # 尝试两种方向（横向和纵向）
        orientations = [
            photo_size,  # 原始方向
            (photo_size[1], photo_size[0])  # 旋转 90 度
        ]

        best_orientation = None
        max_photos = 0

        # 选择最佳方向
        for orientation in orientations:
            # 计算每行和每列可以放置的证件照数量
            photos_per_row = (paper_size[0] - margin) // (orientation[0] + margin)
            photos_per_col = (paper_size[1] - margin) // (orientation[1] + margin)
            total_photos = photos_per_row * photos_per_col

            # 选择可以放置最多证件照的方向
            if total_photos > max_photos:
                max_photos = total_photos
                best_orientation = orientation

        # 使用最佳方向
        photo_size = best_orientation
        if photo_size != orientations[0]:
            photo = photo.rotate(90, expand=True)  # 旋转照片

        # 计算每行和每列可以放置的证件照数量
        photos_per_row = (paper_size[0] - margin) // (photo_size[0] + margin)
        photos_per_col = (paper_size[1] - margin) // (photo_size[1] + margin)

        # 计算总宽度和总高度（包括边距）
        total_width = photos_per_row * photo_size[0] + (photos_per_row - 1) * margin
        total_height = photos_per_col * photo_size[1] + (photos_per_col - 1) * margin

        # 计算起始位置，确保居中
        start_x = (paper_size[0] - total_width) // 2
        start_y = (paper_size[1] - total_height) // 2

        x, y = start_x, start_y  # 初始位置

        # 在相纸上均匀排列证件照
        for row in range(photos_per_col):
            for col in range(photos_per_row):
                # 将证件照粘贴到相纸上
                paper.paste(photo, (x, y))

                # 在每张证件照边缘添加裁剪线，并延伸到相纸边缘
                self.add_crop_lines(draw, x, y, photo_size, paper_size, margin)

                x += photo_size[0] + margin  # 更新 x 坐标

            x = start_x  # 重置 x 坐标
            y += photo_size[1] + margin  # 更新 y 坐标

        return paper
    
    def save_single_photo(self, image, output_folder, size, filename):
        """保存单张证件照"""
        output_path = os.path.join(output_folder, f"{filename}_{size[0]}x{size[1]}.jpg")
        image.save(output_path)
    
    def save_arranged_photo(self, paper, output_folder, paper_size, photo_size, filename):
        """保存排版后的相纸"""
        output_path = os.path.join(output_folder, f"{filename}_paper_{paper_size[0]}x{paper_size[1]}_photo_{photo_size[0]}x{photo_size[1]}.jpg")
        paper.save(output_path)
    
    def make_id_photo(self, image_path, output_folder, size, bg_color, paper_size, margin=50, save_single=True, save_arranged=True):
        """制作证件照并自动排版"""
        try:
            # 打开图像
            image = Image.open(image_path)

            # 调整图像尺寸
            resized_image = self.resize_image(image, size)

            # 替换背景颜色
            final_image = self.change_background_color(resized_image, bg_color)

            # 获取文件名（不含扩展名）
            filename = os.path.splitext(os.path.basename(image_path))[0]

            # 保存单张证件照
            if save_single:
                self.save_single_photo(final_image, output_folder, size, filename)

            # 如果需要保存排版照片
            if save_arranged:
                arranged_paper = self.arrange_photos_on_paper(final_image, paper_size, size, margin)
                self.save_arranged_photo(arranged_paper, output_folder, paper_size, size, filename)
        except Exception as e:
            raise Exception(f"处理图像 {image_path} 时出错: {e}")
    
    def batch_process_images(self, input_folder, output_folder, size, bg_color, paper_size, margin=50, save_single=True, save_arranged=True):
        """批量处理指定文件夹内的照片"""
        # 确保输出文件夹存在
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)

        # 遍历输入文件夹内的所有图像文件
        for filename in os.listdir(input_folder):
            if filename.lower().endswith((".jpg", ".jpeg", ".png")):
                image_path = os.path.join(input_folder, filename)
                try:
                    self.make_id_photo(
                        image_path, output_folder, size, bg_color, 
                        paper_size, margin, save_single, save_arranged
                    )
                except Exception as e:
                    print(f"处理图像 {image_path} 时出错: {e}")

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = PhotoGeneratorApp()
    window.show()
    sys.exit(app.exec())